#include "kernel_operator.h"

// 新增：与SyncAll代码保持一致的同步所需默认大小定义
const int32_t DEFAULT_SYNCALL_NEED_SIZE = 8;

// 原有常量保留，补充注释说明
constexpr int32_t TOTAL_LENGTH = 8 * 2048;                            // 数据总长度
constexpr int32_t USE_CORE_NUM = 8;                                   // 固定使用的核心数量（对应SyncAll的blockNum）
constexpr int32_t BLOCK_LENGTH = TOTAL_LENGTH / USE_CORE_NUM;         // 每个核心处理的数据长度（对应SyncAll的perBlockSize）
constexpr int32_t TILE_NUM = 8;                                       // 每个核心的数据分片数
constexpr int32_t BUFFER_NUM = 2;                                     // 双缓冲设计，提升数据吞吐
constexpr int32_t TILE_LENGTH = BLOCK_LENGTH / TILE_NUM / BUFFER_NUM; // 单个数据片长度（适配双缓冲）
constexpr float EPSILON = 1e-12f;                                      // 数值稳定性保护因子
constexpr int32_t SLOT_STRIDE = 16; 
class KernelLpNormV2 {
public:
    __aicore__ inline KernelLpNormV2() {}

    // 修正：1. 对齐SyncAll的Init参数（添加syncGm，移除冗余workspaceSize）
    //       2. 补充核ID获取（blockIdx），与SyncAll逻辑一致
    __aicore__ inline void Init(GM_ADDR inputGm, GM_ADDR outputGm, 
                                GM_ADDR workGm, GM_ADDR syncGm, 
                                float p, int32_t axis = -1) {
        pipe.Init(); 
        blockIdx = AscendC::GetBlockIdx(); // 新增：获取当前核ID（与SyncAll一致）

        // 1. 全局输入/输出缓冲区初始化（对齐SyncAll的地址偏移逻辑）
        inputGlobal.SetGlobalBuffer(
            (__gm__ float *)inputGm + blockIdx * BLOCK_LENGTH,
            BLOCK_LENGTH
        );
        outputGlobal.SetGlobalBuffer(
            (__gm__ float *)outputGm + blockIdx * BLOCK_LENGTH,
            BLOCK_LENGTH
        );

        // 2. 多核求和缓冲区初始化（使用workGm，大小=核心数×float长度）
        globalSumGlobal.SetGlobalBuffer(
            (__gm__ float *)workGm,
            USE_CORE_NUM * SLOT_STRIDE // 每个核存储1个局部和，共USE_CORE_NUM个元素
        );

        // 3. 同步缓冲区初始化（与SyncAll的syncGlobal逻辑完全对齐）
        syncGlobal.SetGlobalBuffer(
            (__gm__ int32_t *)syncGm,
            USE_CORE_NUM * DEFAULT_SYNCALL_NEED_SIZE // 同步所需总大小
        );

        // 4. 数据队列初始化（保留双缓冲，新增workQueue用于SyncAll）
        pipe.InitBuffer(inQueueInput, BUFFER_NUM, TILE_LENGTH * sizeof(float));
        pipe.InitBuffer(outQueueOutput, BUFFER_NUM, TILE_LENGTH * sizeof(float));
        pipe.InitBuffer(workQueue, 1, USE_CORE_NUM * DEFAULT_SYNCALL_NEED_SIZE * sizeof(int32_t)); // 同步用工作队列

        // 5. 中间计算缓冲区初始化（保留原有逻辑，补充类型强转）
        pipe.InitBuffer(poweredBuf, TILE_LENGTH * sizeof(float));
        pipe.InitBuffer(tmp_p, TILE_LENGTH * sizeof(float));
        pipe.InitBuffer(tmp_n, TILE_LENGTH * sizeof(float));
        pipe.InitBuffer(tmp_base, 1 * sizeof(float));
        pipe.InitBuffer(tmp_epsilon, 1 * sizeof(float));
        pipe.InitBuffer(tmp_exponent, 1 * sizeof(float));
        pipe.InitBuffer(tmp_result, 1 * sizeof(float));
        pipe.InitBuffer(tmp_copy, 1 * sizeof(float)); // 新增：通用临时缓冲区

        // 6. 算法参数初始化
        this->p = p;
        this->axis = axis;
        this->inv_p = 1.0f / p;
        this->pNorm = 0.0f; // 初始化范数结果
    }

    // 核心处理函数（优化同步逻辑，对齐SyncAll的"计算→同步→聚合"流程）
    __aicore__ inline void Process() {
        int32_t loopCount = TILE_NUM  * BUFFER_NUM;
        float localSum = 0.0f;

        // 第一阶段：每个核计算局部p次方和（原有逻辑保留）
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            localSum += ComputePowerSum(i);
        }  

        // 第二阶段：多核同步+全局求和+范数计算（关键优化：对齐SyncAll的SyncAll调用方式）
        ComputeGlobalNorm(localSum);

        // 第三阶段：基于全局范数归一化（原有逻辑保留）
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Normalize(i);
            CopyOut(i);
        }
    }

private:
    // ------------------------------ 数据输入输出函数 ------------------------------
    __aicore__ inline void CopyIn(int32_t progress) {
        AscendC::LocalTensor<float> inputLocal = inQueueInput.AllocTensor<float>();

        AscendC::DataCopy(
            inputLocal, 
            inputGlobal[progress * TILE_LENGTH], 
            TILE_LENGTH
        );
        inQueueInput.EnQue(inputLocal);
    }


    __aicore__ inline void CopyOut(int32_t progress) {
        AscendC::LocalTensor<float> outputLocal = outQueueOutput.DeQue<float>();
        AscendC::DataCopy(
            outputGlobal[progress * TILE_LENGTH], 
            outputLocal, 
            TILE_LENGTH
        );
// 调试信息保留（建议量产时删除）
        // AscendC::PRINTF("[Core %d] CopyOut progress %d, first value = %f\n", 
        //            blockIdx, progress, outputLocal.GetValue(0));
        outQueueOutput.FreeTensor(outputLocal);
    }

    // ------------------------------ 核心计算函数 ------------------------------
    __aicore__ inline float ComputePowerSum(int32_t progress) {
        AscendC::LocalTensor<float> inputLocal = inQueueInput.DeQue<float>();
        AscendC::LocalTensor<float> poweredTensor = poweredBuf.Get<float>();
        AscendC::LocalTensor<float> pTensor = tmp_p.Get<float>();

        // 1. 复制p值到局部张量（适配向量计算）
        AscendC::Duplicate(pTensor, p, TILE_LENGTH);
        // 2. 计算元素绝对值（确保p次方非负）
        AscendC::Abs(poweredTensor, inputLocal, TILE_LENGTH);
        // 3. 计算元素p次方
        AscendC::Power(poweredTensor, poweredTensor, pTensor, TILE_LENGTH);

        // 4. 累加局部和（向量归约）
        float localTileSum = 0.0f;
        for (int32_t i = 0; i < TILE_LENGTH; i++) {
            localTileSum += poweredTensor.GetValue(i);
        }

        inQueueInput.FreeTensor(inputLocal);
        return localTileSum;
    }

    __aicore__ inline void ComputeGlobalNorm(float localSum) {
        // Step1：当前核将局部和写入全局缓冲区
        int32_t slotIdx = blockIdx * SLOT_STRIDE;
        globalSumGlobal.SetValue(slotIdx, localSum);

        AscendC::DataCacheCleanAndInvalid<float,AscendC::CacheLine::SINGLE_CACHE_LINE,AscendC::DcciDst::CACHELINE_OUT>(globalSumGlobal[slotIdx]);

        AscendC::PRINTF("[Core %d] Write localSum = %f to globalSum[%d]\n", 
                   blockIdx, localSum, blockIdx);
        AscendC::PRINTF("[Core %d] actual Writed localSum = %f \n", 
                   blockIdx, globalSumGlobal.GetValue(slotIdx));           

        // Step2：第一次同步（等待所有核完成局部和写入）- 对齐SyncAll逻辑
        AscendC::LocalTensor<int32_t> workLocal = workQueue.AllocTensor<int32_t>();
        AscendC::SyncAll(syncGlobal, workLocal); // 必须传入syncGlobal和workLocal
        workQueue.FreeTensor(workLocal);
        AscendC::PRINTF("[Core %d] First SyncAll completed\n", blockIdx);


        // Step3：仅Core0执行全局求和（避免多核重复计算）
        if (blockIdx == 0) {
            float totalGlobalSum = 0.0f;
            for (int32_t i = 0; i < USE_CORE_NUM; i++) {
                int32_t readIdx = i * SLOT_STRIDE;
                AscendC::DataCacheCleanAndInvalid<float,AscendC::CacheLine::SINGLE_CACHE_LINE,AscendC::DcciDst::CACHELINE_OUT>(globalSumGlobal[readIdx]);

                totalGlobalSum += globalSumGlobal.GetValue(readIdx);
                AscendC::PRINTF("[Core 0] Read globalSum[%d] = %f\n", i, globalSumGlobal.GetValue(readIdx));
            }
            globalSumGlobal.SetValue(0, totalGlobalSum); // 结果写入全局缓冲区第0位
            AscendC::DataCacheCleanAndInvalid<float,AscendC::CacheLine::SINGLE_CACHE_LINE,AscendC::DcciDst::CACHELINE_OUT>(globalSumGlobal[0]);
            AscendC::PRINTF("[Core 0] Global sum computed: %f\n", totalGlobalSum);
        }

        // Step4：第二次同步（等待Core0完成全局求和，确保所有核可见结果）
        workLocal = workQueue.AllocTensor<int32_t>();
        AscendC::SyncAll(syncGlobal, workLocal);
        workQueue.FreeTensor(workLocal);
        AscendC::PRINTF("[Core %d] Second SyncAll completed\n", blockIdx);

        // Step5：所有核读取全局和，计算p范数（(sum + ε)^(1/p)）
        float globalSum = globalSumGlobal.GetValue(0);
        AscendC::LocalTensor<float> baseTensor = tmp_base.Get<float>();
        AscendC::LocalTensor<float> epsilonTensor = tmp_epsilon.Get<float>();
        AscendC::LocalTensor<float> exponentTensor = tmp_exponent.Get<float>();
        AscendC::LocalTensor<float> resultTensor = tmp_result.Get<float>();

        // 数值稳定性处理：sum + ε
        baseTensor.SetValue(0, globalSum);
        epsilonTensor.SetValue(0, EPSILON);
        AscendC::Add(baseTensor, baseTensor, epsilonTensor, 1);

        // 计算p次方根：(sum+ε)^(1/p)
        exponentTensor.SetValue(0, inv_p);
        AscendC::Power(resultTensor, baseTensor, exponentTensor, 1);

        pNorm = resultTensor.GetValue(0);
        AscendC::PRINTF("[Core %d] Final pNorm = %f\n", blockIdx, pNorm);
    }

    // 原有归一化逻辑保留，优化参数传递
    __aicore__ inline void Normalize(int32_t progress) {
        AscendC::LocalTensor<float> inputLocal = inQueueInput.DeQue<float>();
        AscendC::LocalTensor<float> outputLocal = outQueueOutput.AllocTensor<float>();
        AscendC::LocalTensor<float> normTensor = tmp_n.Get<float>();

        // 复制全局范数到局部张量（适配向量除法）
        AscendC::Duplicate(normTensor, pNorm, TILE_LENGTH);
        // 元素级归一化：input / pNorm
        AscendC::Div(outputLocal, inputLocal, normTensor, TILE_LENGTH);

 // 调试信息保留
        // AscendC::PRINTF("[Core %d] Normalize progress %d, first value = %f\n",
        //     blockIdx, progress, outputLocal.GetValue(0));

        outQueueOutput.EnQue(outputLocal);
        inQueueInput.FreeTensor(inputLocal);
    }

    // ------------------------------ 成员变量定义 ------------------------------
    // 1. 算法参数
    float p;             // p范数的阶数
    int32_t axis;        // 计算轴（当前实现为全局轴，可扩展）
    float inv_p;         // 1/p（预计算，减少重复计算）
    float pNorm;         // 全局p范数结果
    int32_t blockIdx;    // 当前核ID（与SyncAll的blockIdx一致）

    // 2. 管道与队列（对齐SyncAll的队列设计）
    AscendC::TPipe pipe;                                      // 计算管道
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueInput;  // 输入数据队列（双缓冲）
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueOutput;// 输出数据队列（双缓冲）
    AscendC::TQue<AscendC::TPosition::VECIN, 1> workQueue;   // 同步用工作队列（与SyncAll一致）

    // 3. 全局缓冲区（对齐SyncAll的GlobalTensor设计）
    AscendC::GlobalTensor<float> inputGlobal;    // 输入全局张量
    AscendC::GlobalTensor<float> outputGlobal;   // 输出全局张量
    AscendC::GlobalTensor<float> globalSumGlobal;// 多核求和全局缓冲区
    AscendC::GlobalTensor<int32_t> syncGlobal;   // 同步用全局缓冲区（与SyncAll一致）

    // 4. 局部计算缓冲区（保留原有设计，位置指定为VECCALC）
    AscendC::TBuf<AscendC::TPosition::VECCALC> poweredBuf;   // 元素p次方存储
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmp_p;        // p值临时缓冲区
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmp_n;        // 范数临时缓冲区
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmp_base;     // 求和结果临时缓冲区
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmp_epsilon;  // EPSILON临时缓冲区
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmp_exponent; // 1/p临时缓冲区
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmp_result;   // 范数计算结果缓冲区
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmp_copy;     // 通用临时缓冲区（备用）
};

// ------------------------------ 核函数定义 ------------------------------
extern "C" __global__ __aicore__ void lp_norm_v2(
    GM_ADDR inputGm, 
    GM_ADDR outputGm, 
    GM_ADDR workGm, 
    GM_ADDR syncGm, 
    float p, 
    int32_t axis
) {
    KernelLpNormV2 op;
    op.Init(inputGm, outputGm, workGm, syncGm, p, axis);
    op.Process();
}

// ------------------------------ 主机端调用接口 ------------------------------
#ifndef ASCENDC_CPU_DEBUG
void lp_norm_v2_do(
    uint32_t blockDim,
    void* stream, 
    uint8_t* inputGm, 
    uint8_t* outputGm, 
    uint8_t* workGm, 
    uint8_t* syncGm, 
    float p, 
    int32_t axis
) {

    lp_norm_v2<<<blockDim, nullptr, stream>>>(inputGm, outputGm, workGm, syncGm, p, axis);
}
#endif




//数据从全局变量读了两次，可优化如下
// ------------------------------ 核心处理函数 ------------------------------
// __aicore__ inline void Process() {
//     int32_t loopCount = TILE_NUM  * this->buffernum;
//     float localSum = 0.0f;

//     // 第一阶段：一次 CopyIn，同时计算局部和并缓存输入
//     for (int32_t i = 0; i < loopCount; i++) {
//         // 从GM拷贝 → 存入 inQueueInput
//         CopyIn(i);

//         // 从输入队列拿数据（读一次）
//         AscendC::LocalTensor<float> inputLocal = inQueueInput.DeQue<float>();

//         // 用于计算局部和
//         localSum += ComputePowerSum(inputLocal);

//         // 复制一份输入缓存起来（供归一化用）
//         AscendC::LocalTensor<float> cacheTensor = cacheQueueInput.AllocTensor<float>();
//         AscendC::DataCopy(cacheTensor, inputLocal, TILE_LENGTH);
//         cacheQueueInput.EnQue(cacheTensor);

//         inQueueInput.FreeTensor(inputLocal);
//     }

//     // 第二阶段：多核同步 + 全局归约 + 范数计算
//     ComputeGlobalNorm(localSum);

//     // 第三阶段：归一化输出（直接用缓存的输入）
//     for (int32_t i = 0; i < loopCount; i++) {
//         AscendC::LocalTensor<float> cacheTensor = cacheQueueInput.DeQue<float>();
//         Normalize(cacheTensor, i);
//         cacheQueueInput.FreeTensor(cacheTensor);
//     }
// }

// // ------------------------------ 改动后的 ComputePowerSum ------------------------------
// __aicore__ inline float ComputePowerSum(AscendC::LocalTensor<float>& inputLocal) {
//     AscendC::LocalTensor<float> poweredTensor = poweredBuf.Get<float>();
//     AscendC::LocalTensor<float> pTensor = tmp_p.Get<float>();

//     AscendC::Duplicate(pTensor, p, TILE_LENGTH);          // p值
//     AscendC::Abs(poweredTensor, inputLocal, TILE_LENGTH); // 绝对值
//     AscendC::Power(poweredTensor, poweredTensor, pTensor, TILE_LENGTH); // p次方

//     // 向量归约求和
//     float localTileSum = 0.0f;
//     for (int32_t i = 0; i < TILE_LENGTH; i++) {
//         localTileSum += poweredTensor.GetValue(i);
//     }
//     return localTileSum;
// }

// // ------------------------------ 改动后的 Normalize ------------------------------
// __aicore__ inline void Normalize(AscendC::LocalTensor<float>& inputLocal, int32_t progress) {
//     AscendC::LocalTensor<float> outputLocal = outQueueOutput.AllocTensor<float>();
//     AscendC::LocalTensor<float> normTensor = tmp_n.Get<float>();

//     AscendC::Duplicate(normTensor, pNorm, TILE_LENGTH); // 范数填充
//     AscendC::Div(outputLocal, inputLocal, normTensor, TILE_LENGTH);

//     outQueueOutput.EnQue(outputLocal);

//     // 输出写回GM
//     CopyOut(progress);
// }

// // ------------------------------ 新增：缓存输入的队列 ------------------------------
AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> cacheQueueInput;
